# --*-- coding: UTF-8 -*-

import h5py
import json
import numpy as np
from utils import image
from bunch import Bunch


class DataLoader():
    def __init__(self, config, mode=0):
        if mode == 0:
            print('training data loading...')
        else:
            print('val data loading...')
        self.config = config
        with h5py.File(self.config.h5_file_path, 'r') as f:

            if mode == 0:
                self.ques = f['ques_train'][:]
                self.img_pos = f['img_pos_train'][:]
                self.answer = f['answers'][:]

            else:
                self.ques = f['ques_val'][:]
                self.img_pos = f['img_pos_val'][:]
                self.answer = f['ans_val'][:]

        data_set = self.get_data_from_json(self.config.json_file_path)

        if mode == 0:
            if self.config.is_have_img_data:
                with h5py.File(config.img_h5_path) as f:
                    self.img_feat_data = f['feat_train'][:]
            self.unique_img = data_set.ix_to_img_train
        else:
            if self.config.is_have_img_data:
                with h5py.File(config.img_h5_path) as f:
                    self.img_feat_data = f['feat_val'][:]
            self.unique_img = data_set.ix_to_img_val

        self.itoa = data_set.ix_to_ans
        self.ques_data_size = len(data_set.ix_to_word)
        self.data_size = len(self.ques)

    def get_data_size(self):
        return self.data_size

    def get_ques_size(self):
        return self.ques_data_size

    def get_ans_size(self):
        return len(self.itoa)

    def get_next_batch(self, epochs_num, batch_size):

        pos = epochs_num * batch_size
        if pos > self.data_size:
            pos = pos % self.data_size

        img = None
        if self.config.is_have_img_data:
            img = np.zeros([batch_size, self.config.fc7_feat_length])
        else:
            img = np.zeros([batch_size, 224, 224, 3])
        ques = np.zeros([batch_size, self.config.max_ques_length], dtype=np.int32)
        answer = np.zeros([batch_size, self.config.ans_vocab_size + 1])

        for idx in range(batch_size):

            ix = (pos + idx) % self.data_size

            ques[idx, :] = self.ques[ix]

            if self.answer[ix] <= 1000:
                answer[idx][self.answer[ix]] = 1.0
            img_idx = self.img_pos[ix]

            if self.config.is_have_img_data:
                img[idx, :] = self.img_feat_data[img_idx]
            else:
                img_path = "%s/%s" % (self.config.img_root, self.unique_img[str(img_idx)])
                img[idx, ...] = image.load_image(img_path)
        return ques, img, answer

    def get_data_from_json(self, path):
        with open(path, 'r') as config_file:
            data_dict = json.load(config_file)

        # convert the dictionary to a namespace using bunch lib
        data = Bunch(data_dict)
        return data


def main():
    conf, _ = config.get_config_from_json('../configs/lstm.json')
    data = DataLoader(conf)

    _, img, _ = data.get_next_batch(0, 10)
    print(img.shape)


if __name__ == '__main__':
    import tensorflow as tf
    from utils import config

    main()